/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "task_builder/task_builder.h"

#include <securec.h>
#include <string>
#include "common/comm_log.h"
#include "common/common_utils.h"
#include "common/comm_error_codes.h"
#include "common/fe_error_code.h"
#include "common/aicore_util_types.h"
#include "common/aicore_util_attr_define.h"
#include "adapter/factory/task_builder_adapter_factory.h"
#include "graph/utils/node_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "util/error_manager/error_manager.h"
#include "runtime/rt_error_codes.h"
#include "runtime/rt_model.h"
#include "runtime/mem.h"

namespace fe {

static const uint32_t FLOWTABLE_ARGS_COUNT = 2;
std::mutex g_task_builder_mutex;
map<rtModel_t, TaskBuilder *> TaskBuilder::task_builder_map_;

TaskBuilder::TaskBuilder() {}

TaskBuilder::~TaskBuilder() {}

Status CreateHandle(ccHandle_t *handle) {
  if (handle == nullptr) {
    REPORT_CM_ERROR("[GenTask][CreatHandle] handle is nullptr");
    return TASK_BUILDER_STATUS_BAD_PARAM;
  }
  void *ccContext = nullptr;
  // alloc ccContext_t
  rtError_t ret;
  ret = rtMallocHost(&ccContext, sizeof(ccContext_t));
  if ((ret != ACL_RT_SUCCESS) || (ccContext == nullptr)) {
    REPORT_CM_ERROR("[GenTask][CreatHandle] alloc handle entity failed");
    return FAILED;
  }
  (void)memset_s(ccContext, sizeof(ccContext_t), 0, sizeof(ccContext_t));
  *handle = static_cast<ccHandle_t>(ccContext);
  return SUCCESS;
}

Status SetStream(ccHandle_t handle, rtStream_t streamId) {
  CM_CHECK_NOTNULL(handle);
  handle->streamId = streamId;
  return SUCCESS;
}

Status DestroyHandle(ccHandle_t *handle) {
  if (handle == nullptr || *handle == nullptr) {
    REPORT_CM_ERROR("[GenTask][DestroyHandle] handle is nullptr");
    return TASK_BUILDER_STATUS_BAD_PARAM;
  }
  ccClearOpMap(*handle);
  rtError_t ret;
  ret = rtFreeHost(*handle);
  if (ret != ACL_RT_SUCCESS) {
    REPORT_CM_ERROR("[GenTask][DestroyHandle] free handler failed");
    return FAILED;
  }
  *handle = nullptr;
  return SUCCESS;
}

Status TaskBuilder::GenerateTask(const ge::Node &node, const ge::RunContext &context,
                                 std::vector<domi::TaskDef> &task_defs) {
  CM_LOGD("TaskBuilder::GenerateTask begin, node name:%s, node type:%s.", node.GetName().c_str(),
          node.GetType().c_str());
  int64_t start_usec_gentask = GetMicroSecondsTime();

  CM_CHECK_NOTNULL(context.model);
  CM_CHECK_NOTNULL(context.stream);
  CM_CHECK_NOTNULL(context.dataMemBase);
  CM_CHECK_NOTNULL(context.weightMemBase);

  ccHandle_t handle = nullptr;
  Status ret = CreateHandle(&handle);
  if (ret != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][CreateHandle][Node %s type %s] CreateHandle failed. ret:0x%X",
                    node.GetName().c_str(), node.GetType().c_str(), ret);
    return FAILED;
  }

  ret = SetStream(handle, context.stream);
  if (ret != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][SetStream][Node %s type %s] SetStream failed. ret:0x%X",
                    node.GetName().c_str(), node.GetType().c_str(), ret);
    (void)DestroyHandle(&handle);
    return FAILED;
  }

  ret = ccSetKernelOpMap(handle);
  if (ret != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][ccSetKernelOpMap][Node %s type %s] ccSetKernelOpMap failed. ret:0x%X",
                    node.GetName().c_str(), node.GetType().c_str(), ret);
    (void)DestroyHandle(&handle);
    return FAILED;
  }

  context_.dataMemSize = context.dataMemSize;
  context_.dataMemBase = context.dataMemBase;
  context_.weightMemSize = context.weightMemSize;
  context_.weightMemBase = context.weightMemBase;
  context_.weightBufferHost = context.weightsBuffer;
  context_.model = context.model;
  context_.stream = context.stream;
  context_.handle = handle;
  task_defs_ = &task_defs;

  // Save current TaskBuilder for runtime callback
  (void)AddTaskBuilder();
  Status status = DoGenerate(node);
  (void)RemoveTaskBuilder();

  (void)DestroyHandle(&handle);

  CM_LOGD("TaskBuilder::GenerateTask end, node name:%s, node type:%s.", node.GetName().c_str(), node.GetType().c_str());

  int64_t end_usec_gentask = GetMicroSecondsTime();
  CM_LOGV("[FE_PERFORMANCE]The time cost of TaskBuilder::GenerateTask is [%ld] micro second.",
          (end_usec_gentask - start_usec_gentask));
  return status;
}

Status TaskBuilder::DoGenerate(const ge::Node &node) {
  // Create TaskBuilderAdapter
  Status status = CreateAdapter(node);
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][DoGenerate][Node %s type %s] Failed to create adapter.",
                    node.GetName().c_str(), node.GetType().c_str());
    return TASK_BUILDER_CREATE_ADAPTER_FAILED;
  }

  bool is_first_node = false;
  bool is_last_node = false;
  (void)ge::AttrUtils::GetBool(node.GetOpDesc(), "is_first_node", is_first_node);
  (void)ge::AttrUtils::GetBool(node.GetOpDesc(), "is_last_node", is_last_node);

  CM_LOGD("TaskBuilder::DoGenerate, node name:%s, node type:%s, first node flag %u, last node flag %u.",
          node.GetName().c_str(), node.GetType().c_str(), is_first_node, is_last_node);
  do {
    rtError_t rtRet = rtSetTaskGenCallback(TaskGenCallback);
    if (rtRet != ACL_RT_SUCCESS) {
      std::map<std::string, std::string> error_key_map;

      // opname,optype,graph_id
      error_key_map[EM_OP_NAME] = node.GetOpDesc()->GetName();
      error_key_map[EM_OP_TYPE] = node.GetOpDesc()->GetType();
      std::string session_graph_id = "";
      (void)ge::AttrUtils::GetStr(node.GetOpDesc(), ge::ATTR_NAME_SESSION_GRAPH_ID, session_graph_id);
      error_key_map[EM_GRAPH_ID] = session_graph_id;
      (void)ErrorManager::GetInstance().ReportErrMessage(EM_GENERATE_TASK_FAILED, error_key_map);

      REPORT_CM_ERROR(
            "[GenTask][DoGenerate][Node %s type %s] rtSetTaskGenCallback failed, ret[0x%X]",
            node.GetName().c_str(), node.GetType().c_str(), rtRet);
      status = TASK_BUILDER_CALL_RT_FAILED;
      break;
    }

    if (is_first_node) {
      rtRet = rtKernelFusionStart(context_.stream);
      if (rtRet != ACL_RT_SUCCESS) {
        REPORT_CM_ERROR(
            "[GenTask][DoGenerate][Node %s type %s]rtKernelFusionStart failed, ret[0x%X]",
            node.GetName().c_str(), node.GetType().c_str(), rtRet);
        status = TASK_BUILDER_CALL_RT_FAILED;
        break;
      }
    }

    // Run TaskBuilderAdapter
    status = RunAdapter();
    if (status != SUCCESS) {
      REPORT_CM_ERROR(
            "[GenTask][DoGenerate][Node %s type %s] Run adapter failed, ret[0x%X]",
            node.GetName().c_str(), node.GetType().c_str(), status);
      status = TASK_BUILDER_RUN_ADAPTER_FAILED;
      break;
    }

    if (is_last_node) {
      rtRet = rtKernelFusionEnd(context_.stream);
      if (rtRet != ACL_RT_SUCCESS) {
        REPORT_CM_ERROR("[GenTask][DoGen][KernelFusion][Node %s type %s] rtKernelFusionEnd failed, ret[0x%X]",
                        node.GetName().c_str(), node.GetType().c_str(), rtRet);
        status = TASK_BUILDER_CALL_RT_FAILED;
        break;
      }
    }
  } while (0);

  return status;
}

Status TaskBuilder::CreateAdapter(const ge::Node &node) {
  ge::OpDescPtr op_desc = node.GetOpDesc();

  int32_t imply_type = -1;
  if (!ge::AttrUtils::GetInt(op_desc, FE_IMPLY_TYPE, imply_type)) {
    REPORT_CM_ERROR("[GenTask][CreateAdapter][Node %s type %s] Get op imply_type failed, imply_type[%u]]",
                        op_desc->GetName().c_str(), op_desc->GetType().c_str(), imply_type);
    return FAILED;
  }
  auto op_impl_type = static_cast<OpImplType>(imply_type);

  adapter_ = TaskBuilderAdapterFactory::Instance().CreateTaskBuilderAdapter(op_impl_type, node, context_);
  CM_CHECK_NOTNULL(adapter_);

  CM_LOGD("Create TaskBuilderAdapter success. OpName:%s, OpType:%s",
          op_desc->GetName().c_str(), op_desc->GetType().c_str());

  Status status = adapter_->Init();
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][CreateAdapter][Node %s type %s] Init TaskBuilderAdapter failed.",
                        op_desc->GetName().c_str(), op_desc->GetType().c_str());
    return status;
  }

  return SUCCESS;
}

Status TaskBuilder::RunAdapter() {
  ge::ConstOpDescPtr op_desc = adapter_->GetOpDesc();
  CM_LOGD("start run TaskBuilderAdapter, OpName:%s, OpType:%s", op_desc->GetName().c_str(), op_desc->GetType().c_str());

  orig_op_indexes_.clear();
  orig_op_indexes_.push_back(static_cast<uint32_t>(op_desc->GetId()));
  CM_LOGD("Push back node(%s)'s op_desc id(%d) into orig_op_indexes success.", op_desc->GetName().c_str(),
          orig_op_indexes_.back());

  if (context_.handle == nullptr) {
    REPORT_CM_ERROR("[GenTask][RunAdapter][Node %s type %s] handle is nullptr!",
                    op_desc->GetName().c_str(), op_desc->GetType().c_str());
    return FAILED;
  }
  context_.handle->opIndex = static_cast<uint32_t>(op_desc->GetId());

  Status status = adapter_->Run();
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][RunAdapter][Node %s type %s] Fail to run TaskBuilderAdapter.",
                    op_desc->GetName().c_str(), op_desc->GetType().c_str());
    return status;
  }

  CM_LOGD("TaskBuilderAdapter run success, OpName:%s, OpType:%s",
          op_desc->GetName().c_str(), op_desc->GetType().c_str());

  return SUCCESS;
}

void TaskBuilder::AddTaskBuilder() {
  std::lock_guard<std::mutex> lock_guard(g_task_builder_mutex);
  task_builder_map_[context_.model] = this;
}

void TaskBuilder::RemoveTaskBuilder() {
  std::lock_guard<std::mutex> lock_guard(g_task_builder_mutex);
  (void)task_builder_map_.erase(context_.model);
}

TaskBuilder *TaskBuilder::GetTaskBuilderByModel(const rtModel_t &model) {
  std::lock_guard<std::mutex> lock_guard(g_task_builder_mutex);
  map<rtModel_t, TaskBuilder *>::iterator iter = task_builder_map_.find(model);
  if (iter == task_builder_map_.end()) {
    REPORT_CM_ERROR("[GenTask][GetTaskBuilderByModel] TaskBuilder not found.");
    return nullptr;
  }
  return iter->second;
}

rtError_t TaskBuilder::TaskGenCallback(rtModel_t model, rtTaskInfo_t *task_info) {
  if (model == nullptr || task_info == nullptr) {
    REPORT_CM_ERROR("[GenTask][TaskGenCallback] model or task_info is nullptr.");
    return ACL_ERROR_RT_PARAM_INVALID;
  }
  TaskBuilder *builder = GetTaskBuilderByModel(model);
  if (builder == nullptr || builder->task_defs_ == nullptr) {
    REPORT_CM_ERROR("[GenTask][TaskGenCallback] TaskBuilder is invalid.");
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  domi::TaskDef task_def = {};
  task_def.set_type(task_info->type);
  task_def.set_stream_id(task_info->streamID);

  CM_LOGD("TaskInfo: task type = %d", task_info->type);

  if (task_info->type == (uint32_t)RT_MODEL_TASK_KERNEL) {
    rtError_t rtRet = GenKernelTask(*builder, *task_info, task_def);
    if (rtRet != ACL_RT_SUCCESS) {
      REPORT_CM_ERROR("[GenTask][TaskGenCallback] GenKernelTask failed. ret:0x%X", rtRet);
      return rtRet;
    }
  } else if (task_info->type == (uint32_t)RT_MODEL_TASK_EVENT_RECORD ||
             task_info->type == (uint32_t)RT_MODEL_TASK_EVENT_WAIT) {
    task_def.set_event_id(task_info->u.eventTask.eventID);
  } else if (task_info->type == (uint32_t)RT_MODEL_TASK_KERNEL_EX) {
    domi::KernelExDef *kernel_def_ex = task_def.mutable_kernel_ex();
    CM_LOGI("KernelExDef: args_size = %d", task_info->u.kernelTaskEx.argsSize);

    if (builder->orig_op_indexes_.empty()) {
      REPORT_CM_ERROR("[GenTask][TaskGenCallback] The value[builder->orig_op_indexes_.size()] is 0.");
      return ACL_ERROR_RT_PARAM_INVALID;
    }

    kernel_def_ex->set_flags(task_info->u.kernelTaskEx.flags);
    kernel_def_ex->set_op_index(builder->orig_op_indexes_[0]);
    kernel_def_ex->set_args_size(task_info->u.kernelTaskEx.argsSize);
    kernel_def_ex->set_args(task_info->u.kernelTaskEx.args, task_info->u.kernelTaskEx.argsSize);
  } else if(task_info->type == (uint32_t)RT_MODEL_TASK_ALL_KERNEL) {
    rtError_t rtRet = GenAllKernelTask(*builder, *task_info, task_def);
    if (rtRet != ACL_RT_SUCCESS) {
      REPORT_CM_ERROR("[GenTask][TaskGenCallback] GenAllKernelTask failed failed. ret:0x%X", rtRet);
      return rtRet;
    }
  } else {
    // do nothing.
  }

  // Collect TaskDef
  builder->task_defs_->push_back(task_def);

  return ACL_RT_SUCCESS;
}

rtError_t TaskBuilder::GenAllKernelTask(const TaskBuilder &builder, const rtTaskInfo_t &task_info,
                                        domi::TaskDef &task_def) {
  const rtAllKernelTaskInfo_t &kernel_task = task_info.u.allKernelTask;

  CM_LOGD("AllKernelTaskInfo: dev_func = %s, block_dim = %d, args_size = %d",
          kernel_task.devfunc, kernel_task.blockDim, kernel_task.argsSize);

  domi::KernelDefWithHandle *kernel_def_with_handle = task_def.mutable_kernel_with_handle();
  if (kernel_def_with_handle == nullptr) {
    CM_LOGD("Invalid ptr initialize, taskInfo: task type = %d", task_info.type);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  domi::KernelContext *kernel_context = kernel_def_with_handle->mutable_context();
  TaskBuilderAdapterPtr adapter = builder.adapter_;
  // Set kernel_name in kernel_def_with_handle
  if (adapter != nullptr) {
    ge::ConstOpDescPtr op_desc = adapter->GetOpDesc();
    string first_kernel_name;
    if (ge::AttrUtils::GetStr(op_desc, ATTR_NAME_KERNEL_LIST_FIRST_NAME, first_kernel_name))  {
      kernel_def_with_handle->set_node_info(first_kernel_name);

      string attr_key_kernel_name = op_desc->GetName() + "_kernelname";
      string attr_kernel_name;
      (void)ge::AttrUtils::GetStr(op_desc, attr_key_kernel_name, attr_kernel_name);
      kernel_def_with_handle->set_original_kernel_key(attr_kernel_name);
    } else {
      REPORT_CM_ERROR("[GenTask][GenAllKernelTask] No kernel list name in op desc.");
      return ACL_ERROR_RT_PARAM_INVALID;
    }
  }

  kernel_def_with_handle->set_block_dim(kernel_task.blockDim);
  kernel_def_with_handle->set_args_size(kernel_task.argsSize);
  kernel_def_with_handle->set_args(kernel_task.args, kernel_task.argsSize);

  ccOpContext op_context;
  Status ret = ccGetKernelContext(builder.context_.stream, op_context);
  if (ret != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][GenAllKernelTask] ccGetKernelContext failed, ret[0x%X]", ret);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  // Fill kernel_context
  Status status = FillKernelContextMultiKernel(builder.orig_op_indexes_, kernel_task, op_context, *kernel_context);
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][GenAllKernelTask] FillKernelContext failed, ret[0x%X]", status);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  return ACL_RT_SUCCESS;
}

rtError_t TaskBuilder::GenKernelTask(const TaskBuilder &builder, const rtTaskInfo_t &task_info,
                                     domi::TaskDef &task_def) {
  const rtKernelTaskInfo_t &kernel_task = task_info.u.kernelTask;

  CM_LOGD("KernelTaskInfo: stub_func = %s, block_dim = %d, args_size = %d", kernel_task.stubFunc, kernel_task.blockDim,
          kernel_task.argsSize);

  domi::KernelDef *kernel_def = task_def.mutable_kernel();
  if (kernel_def == nullptr) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] kernel_def is nullptr.");
    return ACL_ERROR_RT_PARAM_INVALID;
  }
  domi::KernelContext *kernel_context = kernel_def->mutable_context();
  if (kernel_context == nullptr) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] kernel_context is nullptr.");
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  // Set kernel_name in kernel_def
  TaskBuilderAdapterPtr adapter = builder.adapter_;
  if (adapter != nullptr) {
    ge::ConstOpDescPtr op_desc = adapter->GetOpDesc();
    string attr_key_kernel_name = op_desc->GetName() + "_kernelname";
    string attr_val_kernel_name;
    (void)ge::AttrUtils::GetStr(op_desc, attr_key_kernel_name, attr_val_kernel_name);

    kernel_def->set_kernel_name(attr_val_kernel_name);
    CM_LOGD("Set kernel_name:%s for the kernel_def of node:%s success.", attr_val_kernel_name.c_str(),
            op_desc->GetName().c_str());
  }

  if (kernel_task.stubFunc == nullptr) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] kernelTask.stub_func is nullptr.");
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  string stub_func_name(kernel_task.stubFunc);
  kernel_def->set_stub_func(stub_func_name);
  kernel_def->set_block_dim(kernel_task.blockDim);

  if (kernel_task.smDesc != nullptr) {
    CM_LOGD("L2 sm_desc is not nullptr.");
    kernel_def->set_sm_desc(kernel_task.smDesc, sizeof(rtSmDesc_t));
  }

  kernel_def->set_args_size(kernel_task.argsSize);
  kernel_def->set_args(kernel_task.args, kernel_task.argsSize);

  ccOpContext op_context;
  Status ret = ccGetKernelContext(builder.context_.stream, op_context);
  if (ret != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] ccGetKernelContext failed, ret[0x%X]", ret);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  // Fill kernel_context
  Status status = FillKernelContext(builder.orig_op_indexes_, kernel_task, op_context, *kernel_context);
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] FillKernelContext failed, ret[0x%X]", status);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  // Fill Flowtable
  status = FillFlowtable(*kernel_context, kernel_task, *kernel_def);
  if (status != SUCCESS) {
    REPORT_CM_ERROR("[GenTask][GenKernelTask] FillFlowtable failed, ret[0x%X]", status);
    return ACL_ERROR_RT_PARAM_INVALID;
  }

  return ACL_RT_SUCCESS;
}

Status TaskBuilder::FillKernelContext(const vector<uint32_t> &orig_op_indexes, const rtKernelTaskInfo_t &kernel_task,
                                      const ccOpContext &op_context, domi::KernelContext &kernel_context) {
  for (size_t i = 0; i < orig_op_indexes.size(); i++) {
    kernel_context.add_origin_op_index(orig_op_indexes[i]);
  }
  kernel_context.set_kernel_type(op_context.kernelType);
  if (op_context.kernelType == CCE_AI_CORE || op_context.kernelType == CCE_AI_CPU) {
    // CCE OP
    CM_LOGI("KernelContext: kernel_type = %d, op_id = %d, kernel_func_id = %d, argsCount = %d",
            op_context.kernelType, op_context.opId, op_context.kernelFuncId, kernel_task.argsCount);

    kernel_context.set_op_id(op_context.opId);
    kernel_context.set_kernel_func_id(op_context.kernelFuncId);
    kernel_context.set_is_flowtable(op_context.isFlowtable);
    kernel_context.set_args_count(kernel_task.argsCount);
    kernel_context.set_args_offset(kernel_task.argsOffset, kernel_task.argsCount * sizeof(uint16_t));
  } else {
    if (orig_op_indexes.empty()) {
      REPORT_CM_ERROR("[GenTask][FillKernelContext] Kernel Context is empty");
      return FAILED;
    }
    kernel_context.set_is_flowtable(false);
    kernel_context.set_op_index(orig_op_indexes[0]);
    kernel_context.set_args_count(kernel_task.argsCount);
    kernel_context.set_args_offset(kernel_task.argsOffset, kernel_task.argsCount * sizeof(uint16_t));
  }

  return SUCCESS;
}

Status TaskBuilder::FillFlowtable(const domi::KernelContext &kernel_context, const rtKernelTaskInfo_t &kernel_task,
                                  domi::KernelDef &kernel_def) {
  if (kernel_context.is_flowtable()) {
    if (kernel_task.argsCount < FLOWTABLE_ARGS_COUNT) {
      REPORT_CM_ERROR("[GenTask][FillFlowtable] kernelTask.args_count < %u", FLOWTABLE_ARGS_COUNT);
      return FAILED;
    }

    // according to argsOffset to obtain flowtable data
    // args data : |...some size info...|...flowtable addr...|...flowtable
    // size...|...
    //             ↓                    ↓                    ↓
    //             args_addr             args_offset[0]        args_offset[1]
    uint32_t flowtable_size = *reinterpret_cast<uint32_t *>(kernel_task.args + kernel_task.argsOffset[1]);
    if (flowtable_size <= 0) {
      REPORT_CM_ERROR("[GenTask][FillFlowtable] The value[flowtable_size] <= 0");
      return FAILED;
    }

    uint64_t addr = *reinterpret_cast<uint64_t *>(kernel_task.args + kernel_task.argsOffset[0]);
    kernel_def.set_flowtable(reinterpret_cast<void *>((uintptr_t)addr), flowtable_size);

    CM_LOGI("FlowtableInfo: Flowtable data size = %d", flowtable_size);
  }

  return SUCCESS;
}

Status TaskBuilder::FillKernelContextMultiKernel(const vector<uint32_t> &orig_op_indexes,
                                                 const rtAllKernelTaskInfo_t &kernel_task,
                                                 const ccOpContext &op_context,
                                                 domi::KernelContext &kernel_context) {
  for (size_t i = 0; i < orig_op_indexes.size(); i++) {
    kernel_context.add_origin_op_index(orig_op_indexes[i]);
  }
  kernel_context.set_kernel_type(op_context.kernelType);
  if (op_context.kernelType == CCE_AI_CORE || op_context.kernelType == CCE_AI_CPU) {
    // CCE OP
    CM_LOGI("KernelContext: kernel_type = %d, op_id = %d, kernel_func_id = %d, argsCount = %d",
            op_context.kernelType, op_context.opId, op_context.kernelFuncId, kernel_task.argsCount);

    kernel_context.set_op_id(op_context.opId);
    kernel_context.set_kernel_func_id(op_context.kernelFuncId);
    kernel_context.set_is_flowtable(op_context.isFlowtable);
    kernel_context.set_args_count(kernel_task.argsCount);
    kernel_context.set_args_offset(kernel_task.argsOffset, kernel_task.argsCount * sizeof(uint16_t));
  } else {
    if (orig_op_indexes.empty()) {
      REPORT_CM_ERROR("[GenTask][FillKernelContext] Kernel Context is empty.");
      return FAILED;
    }
    kernel_context.set_is_flowtable(false);
    kernel_context.set_op_index(orig_op_indexes[0]);
    kernel_context.set_args_count(kernel_task.argsCount);
    kernel_context.set_args_offset(kernel_task.argsOffset, kernel_task.argsCount * sizeof(uint16_t));
  }

  return SUCCESS;
}
}  // namespace fe
